#!/usr/bin/env python3
"""
说话人管理器模块
负责说话人识别、命名、颜色分配和会话管理
"""

import time
from typing import Any, Optional

import numpy as np

# 导入配置
from speech_convert.config.funasr_config import CONFIG


class SpeakerManager:
    """说话人管理器 - 处理说话人识别、命名、颜色分配等"""

    def __init__(self):
        self.speaker_map: dict[str, dict[str, Any]] = {}  # speaker_id -> speaker_info
        self.speaker_embeddings: dict[str, Any] = {}  # speaker_id -> embedding_vector
        self.session_speakers: dict[str, list[str]] = {}  # session_id -> [speaker_ids]
        self.message_counter = 0
        self.session_start_time = time.time()
        # 说话人相似度阈值 - 从配置文件读取
        self.similarity_threshold = CONFIG.speaker.SPEAKER_SIMILARITY_THRESHOLD
        self.speaker_counter = 0  # 说话人计数器

    def _cosine_similarity(self, vec1: Any, vec2: Any) -> float:
        """计算余弦相似度"""
        try:
            # 转换为numpy数组并确保是一维向量
            arr1 = np.array(vec1).flatten()
            arr2 = np.array(vec2).flatten()

            # 检查维度是否匹配
            if arr1.shape[0] != arr2.shape[0]:
                print(f"⚠️ 向量维度不匹配: {arr1.shape} vs {arr2.shape}")
                return 0.0

            # 计算余弦相似度
            dot_product = np.dot(arr1, arr2)
            norm1 = np.linalg.norm(arr1)
            norm2 = np.linalg.norm(arr2)

            if norm1 == 0 or norm2 == 0:
                print("⚠️ 零向量无法计算相似度")
                return 0.0

            similarity = dot_product / (norm1 * norm2)
            return float(similarity)
        except Exception as e:
            print(f"⚠️ 计算相似度错误: {e}")
            print(f"   vec1 shape: {getattr(vec1, 'shape', 'unknown')}")
            print(f"   vec2 shape: {getattr(vec2, 'shape', 'unknown')}")
            return 0.0

    def find_or_create_speaker(self, embedding: Any, confidence: float) -> str:
        """查找相似说话人或创建新说话人"""
        try:
            # 转换为numpy数组
            current_embedding = np.array(embedding)

            # 查找最相似的说话人
            best_speaker_id = None
            best_similarity = 0.0

            for speaker_id, stored_embedding in self.speaker_embeddings.items():
                similarity = self._cosine_similarity(
                    current_embedding, stored_embedding
                )
                if similarity > best_similarity:
                    best_similarity = similarity
                    best_speaker_id = speaker_id

            # 如果找到相似度高于阈值的说话人，则使用该ID
            if best_speaker_id and best_similarity >= self.similarity_threshold:
                # 更新置信度
                self.speaker_map[best_speaker_id]["confidence"] = max(
                    self.speaker_map[best_speaker_id]["confidence"], confidence
                )
                self.speaker_map[best_speaker_id]["is_new_speaker"] = False
                print(
                    f"🎭 识别已知说话人: {self.speaker_map[best_speaker_id]['name']} "
                    f"(相似度: {best_similarity:.3f})"
                )
                return best_speaker_id

            # 创建新说话人
            self.speaker_counter += 1
            new_speaker_id = f"speaker_{self.speaker_counter:03d}"

            # 先注册新说话人
            self.register(new_speaker_id, confidence)

            # 然后存储嵌入向量
            self.speaker_embeddings[new_speaker_id] = current_embedding

            return new_speaker_id

        except Exception as e:
            print(f"⚠️ 说话人查找/创建错误: {e}")
            # 降级处理：使用简单的hash方式
            embedding_str = str(embedding)
            fallback_id = f"speaker_{abs(hash(embedding_str)) % 1000:03d}"
            self.register(fallback_id, confidence)
            return fallback_id

    def get_session_id(self) -> str:
        """获取当前会话ID"""
        timestamp = int(self.session_start_time)
        return f"session_{timestamp}"

    def get_message_id(self) -> str:
        """生成消息ID"""
        self.message_counter += 1
        return f"msg_{self.message_counter:06d}"

    def register(self, speaker_id: str, confidence: float = 0.0) -> dict[str, Any]:
        """注册新说话人或获取已存在的说话人信息"""
        if speaker_id not in self.speaker_map:
            # 分配说话人编号和颜色
            speaker_count = len(self.speaker_map)
            speaker_index = speaker_count % len(CONFIG.speaker.SPEAKER_COLORS)

            speaker_info = {
                "speaker_id": speaker_id,
                "name": f"{CONFIG.speaker.SPEAKER_NAME_PREFIX}{CONFIG.speaker.SPEAKER_NAME_LETTERS[speaker_index]}",
                "color": CONFIG.speaker.SPEAKER_COLORS[speaker_index],
                "confidence": confidence,
                "is_new_speaker": True,
                "first_seen": time.time(),
            }

            self.speaker_map[speaker_id] = speaker_info

            # 记录到当前会话
            session_id = self.get_session_id()
            if session_id not in self.session_speakers:
                self.session_speakers[session_id] = []
            self.session_speakers[session_id].append(speaker_id)

            print(
                f"🎭 注册新说话人: {speaker_info['name']} ({speaker_id}) 颜色:{speaker_info['color']}"
            )
        else:
            # 更新置信度，标记为非新说话人
            self.speaker_map[speaker_id]["confidence"] = confidence
            self.speaker_map[speaker_id]["is_new_speaker"] = False

        return self.speaker_map[speaker_id].copy()

    def get_speaker_info(self, speaker_id: str) -> dict[str, Any]:
        """获取说话人信息"""
        return self.speaker_map.get(speaker_id, {})

    def get_session_speakers(self, session_id: Optional[str] = None) -> list[str]:
        """获取会话中的所有说话人"""
        if session_id is None:
            session_id = self.get_session_id()
        return self.session_speakers.get(session_id, [])

    def reset_session(self):
        """重置会话"""
        self.session_start_time = time.time()
        self.message_counter = 0
        # 清空说话人缓存，开始新会话
        self.speaker_map.clear()
        self.speaker_embeddings.clear()
        self.speaker_counter = 0
        print(f"🔄 会话已重置: {self.get_session_id()}")

    def get_statistics(self) -> dict[str, Any]:
        """获取统计信息"""
        return {
            "total_speakers": len(self.speaker_map),
            "current_session": self.get_session_id(),
            "message_count": self.message_counter,
            "session_duration": time.time() - self.session_start_time,
            "similarity_threshold": self.similarity_threshold,
        }
